package org.study.shardingsphere.boot.algorithm;

import cn.hutool.core.date.DateUtil;
import com.google.common.collect.Range;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 自定义符合分片算法：按照reconciliationDay-年月+车牌hash分片
 * 精准分片：=、in
 * 范围分片：>、>=、<、<=、between
 *
 * @author lenovo
 */
@Slf4j
public class ReconComplexShardingAlgorithm implements ComplexKeysShardingAlgorithm {
    private final static String RECON_DAY = "reconciliation_day";
    private final static String PLATE_NUMBER = "plate_number";

    @Override
    public Collection<String> doSharding(Collection availableTargetNames, ComplexKeysShardingValue shardingValue) {
        //log.info("精确分片-按照对账日+车牌号分表-availableTargetNames={},shardingValue={}", availableTargetNames, shardingValue);
        // 获取年月,车牌hash的桶两个维度
        SortedSet<Integer> availableReconDaySet = new TreeSet<>();
        SortedSet<Integer> availablePlateNumberHashSet = new TreeSet<>();
        addConfigYearMonthAndPlateNumber(availableTargetNames, availableReconDaySet, availablePlateNumberHashSet);

        // 存储条件中的年月信息
        SortedSet<Integer> yearMonthSet = new TreeSet<>();
        // 存储条件中的车牌哈希桶
        SortedSet<Integer> plateNumberHashSet = new TreeSet<>();

        // 范围条件
        if (!CollectionUtils.isEmpty(shardingValue.getColumnNameAndRangeValuesMap())) {
            Range<Date> range = (Range<Date>) shardingValue.getColumnNameAndRangeValuesMap().get(RECON_DAY);
            // 查询的时间范围
            int lower = range.hasLowerBound() ? getYearMonth(range.lowerEndpoint()) : 197001;
            int upper = range.hasUpperBound() ? getYearMonth(range.upperEndpoint()) : 999912;
            SortedSet<Integer> rangeTables = addRangeTables(Math.max(availableReconDaySet.first(), lower), Math.min(availableReconDaySet.last(), upper));
            yearMonthSet.addAll(rangeTables);
        }

        // 精准条件
        if (!CollectionUtils.isEmpty(shardingValue.getColumnNameAndShardingValuesMap())) {
            Collection<Date> reconDayList = (Collection<Date>) shardingValue.getColumnNameAndShardingValuesMap().get(RECON_DAY);
            Collection<String> plateNumberList = (Collection<String>) shardingValue.getColumnNameAndShardingValuesMap().get(PLATE_NUMBER);

            // 获取条件中的年月
            Set<Integer> yearMonthShardingSet = CollectionUtils.isEmpty(reconDayList) ?
                    new HashSet<>() :
                    reconDayList.stream().map(reconDay -> getYearMonth(reconDay)).collect(Collectors.toSet());
            yearMonthSet.addAll(yearMonthShardingSet);

            // 获取条件中的车牌哈希桶
            Set<Integer> plateNumberShardingHashSet = CollectionUtils.isEmpty(plateNumberList) ?
                    new HashSet<>() :
                    plateNumberList.stream().map(plateNumber -> Math.abs(plateNumber.hashCode()) % availablePlateNumberHashSet.size()).collect(Collectors.toSet());
            plateNumberHashSet.addAll(plateNumberShardingHashSet);
        }

        Set<String> realTables = getRealTables(shardingValue.getLogicTableName(),
                // 条件中没有年月信息,使用配置中的年月
                !CollectionUtils.isEmpty(yearMonthSet) ? yearMonthSet : availableReconDaySet,
                // 条件中没有车牌哈希桶信息,使用配置中车牌哈希桶
                !CollectionUtils.isEmpty(plateNumberHashSet) ? plateNumberHashSet : availablePlateNumberHashSet);
        //log.info("精确分片-按照对账日+车牌号分表-命中表={}", realTables);
        return CollectionUtils.isEmpty(realTables) ? availableTargetNames : realTables;
    }

    /**
     * 拆分年月,车牌hash的桶两个维度
     *
     * @param availableTargetNames
     * @param availableReconDaySet
     * @param availablePlateNumberHashSet
     */
    private void addConfigYearMonthAndPlateNumber(Collection<String> availableTargetNames, SortedSet<Integer> availableReconDaySet, SortedSet<Integer> availablePlateNumberHashSet) {
        for (String targetName : availableTargetNames) {
            //String targetName = (String) availableTargetName;
            // 大概格式:xxx_月份_hash桶号
            String[] targetNameSplit = targetName.split("_");
            // 倒数第二个字段
            availableReconDaySet.add(Integer.valueOf(targetNameSplit[targetNameSplit.length - 2]));
            // 倒数第一个字段
            availablePlateNumberHashSet.add(Integer.valueOf(targetNameSplit[targetNameSplit.length - 1]));
        }
    }

    /**
     * 获取实际表
     *
     * @param logicTableName
     * @param yearMonthSet
     * @param plateNumberHashSet
     * @return
     */
    private Set<String> getRealTables(String logicTableName, Set<Integer> yearMonthSet, Set<Integer> plateNumberHashSet) {
        Set<String> tables = new HashSet<>();
        for (Integer yearMonth : yearMonthSet) {
            for (Integer num : plateNumberHashSet) {
                tables.add(logicTableName + "_" + yearMonth + "_" + num);
            }
        }
        return tables;
    }

    /**
     * 获取年月
     *
     * @param reconDay
     * @return
     */
    private Integer getYearMonth(Date reconDay) {
        //LocalDateTime localDateTime = LocalDateTimeUtils.date2LocalDateTime(reconDay);
        //String format = LocalDateTimeUtils.format(localDateTime, LocalDateTimeUtils.YYYYMM_FORMAT);
        String format = DateUtil.format(reconDay, "yyyyMM");
        return Integer.parseInt(format);
    }

    /**
     * 获取返回时间内的年月
     *
     * @param lowerYearMonth
     * @param upperYearMonth
     */
    private SortedSet<Integer> addRangeTables(Integer lowerYearMonth, Integer upperYearMonth) {
        SortedSet<Integer> yearMonths = new TreeSet<>();
        // 将区间内的年月获取处理
        yearMonths.add(lowerYearMonth);
        for (int i = lowerYearMonth; i <= upperYearMonth; i++) {
            yearMonths.add(i);
        }
        return yearMonths;
    }

    @Override
    public void init(Properties properties) {
    }
}
